import numpy as np
import tensorflow as tf
import util
batch_size=128
def encoder(input):
    with tf.variable_scope('encoder'):
        with tf.variable_scope('layer_1'):
            output1=scale3(input)
            output2=scale5(input)
            output3=scale7(input)
            output=tf.concat([output1,output2,output3],3)
            print(('merge' + str(output.get_shape().as_list())))
            output=util.conv2d(output,1,1)
            output = tf.nn.relu(output)
            output=tf.reshape(output,[-1,1024])

        with tf.variable_scope('layer_2'):
            output=util.fullyConnect(output,np.load('ZIDCT.npy').T)
            print(('dct' + str(output.get_shape().as_list())))
    return output
def scale3(input):
    with tf.variable_scope('scale3'):
        with tf.variable_scope('layer_1'):
            output=input
        with tf.variable_scope('layer_2'):
            output=util.conv2d(output,3,64)
            output = tf.nn.relu(output)
            print('scale3' + str(output.get_shape().as_list()))
        with tf.variable_scope('layer_3'):
            output=util.conv2d(output,3,256)
            output = tf.nn.relu(output)
            print('scale3' + str(output.get_shape().as_list()))
        with tf.variable_scope('layer_4'):
            output=util.conv2d(output,3,512)
            output = tf.nn.relu(output)
            print('scale3' + str(output.get_shape().as_list()))
    return output
def scale5(input):
    with tf.variable_scope('scale5'):
        with tf.variable_scope('layer_1'):
            output=input
        with tf.variable_scope('layer_2'):
            output=util.conv2d(output,5,32)
            output = tf.nn.relu(output)
            print('scale5' + str(output.get_shape().as_list()))
        with tf.variable_scope('layer_3'):
            output=util.conv2d(output,5,64)
            output = tf.nn.relu(output)
            print('scale5' + str(output.get_shape().as_list()))
        with tf.variable_scope('layer_4'):
            output=util.conv2d(output,5,128)
            output = tf.nn.relu(output)
            print('scale5' + str(output.get_shape().as_list()))
        with tf.variable_scope('layer_5'):
            output=util.conv2d(output,5,256)
            output = tf.nn.relu(output)
            print('scale5' + str(output.get_shape().as_list()))
    return output
def scale7(input):
    with tf.variable_scope('scale7'):
        with tf.variable_scope('layer_1'):
            output=input
        with tf.variable_scope('layer_2'):
            output=util.conv2d(output,7,64)
            output = tf.nn.relu(output)
            print('scale7' + str(output.get_shape().as_list()))
        with tf.variable_scope('layer_3'):
            output=util.conv2d(output,7,128)
            output = tf.nn.relu(output)
            print('scale7' + str(output.get_shape().as_list()))
        with tf.variable_scope('layer_4'):
            output=util.conv2d(output,7,256)
            output = tf.nn.relu(output)
            print('scale7' + str(output.get_shape().as_list()))
        with tf.variable_scope('layer_5'):
            output=util.conv2d(output,7,512)
            output = tf.nn.relu(output)
            print('scale7' + str(output.get_shape().as_list()))
    return output
def decoder(input):
    with tf.variable_scope('decoder'):
        with tf.variable_scope('layer_1'):
            output = util.fullyConnect(input, np.load('ZIDCT.npy'))
            print(output.get_shape().as_list())
        with tf.variable_scope('layer_2'):
            output=tf.reshape(output,[-1,32,32,1])
            output1=rescale3(output)
            output2=rescale5(output)
            output3=rescale7(output)
            output = tf.concat([output1, output2, output3], 3)
            print(('remerge' + str(output.get_shape().as_list())))
            output = util.conv2d(output, 1, 1)
            output = tf.nn.relu(output)
            print(('output' + str(output.get_shape().as_list())))
    return output

def rescale3(input):
    with tf.variable_scope('rescale3'):
        with tf.variable_scope('layer_1'):
            output=util.conv2d(input,3,512)
            output=tf.nn.relu(output)
            print('rescale3' + str(output.get_shape().as_list()))
        with tf.variable_scope('layer_2'):
            output=util.conv2d(output,3,256)
            output=tf.nn.relu(output)
            print('rescale3' + str(output.get_shape().as_list()))
        with tf.variable_scope('layer_3'):
            output=util.conv2d(output,3,64)
            output=tf.nn.relu(output)
            print('rescale3' + str(output.get_shape().as_list()))
    return output
def rescale5(input):
    with tf.variable_scope('rescale5'):
        with tf.variable_scope('layer_1'):
            output=util.conv2d(input,5,256)
            output=tf.nn.relu(output)
            print('rescale5' + str(output.get_shape().as_list()))
        with tf.variable_scope('layer_2'):
            output=util.conv2d(output,5,128)
            output=tf.nn.relu(output)
            print('rescale5' + str(output.get_shape().as_list()))
        with tf.variable_scope('layer_3'):
            output=util.conv2d(output,5,64)
            output=tf.nn.relu(output)
            print('rescale5' + str(output.get_shape().as_list()))
        with tf.variable_scope('layer_4'):
            output=util.conv2d(output,5,32)
            output=tf.nn.relu(output)
            print('rescale5' + str(output.get_shape().as_list()))
    return output
def rescale7(input):
    with tf.variable_scope('rescale7'):
        with tf.variable_scope('layer_1'):
            output=util.conv2d(input,7,512)
            output=tf.nn.relu(output)
            print('rescale7' + str(output.get_shape().as_list()))
        with tf.variable_scope('layer_2'):
            output=util.conv2d(output,7,256)
            output=tf.nn.relu(output)
            print('rescale7' + str(output.get_shape().as_list()))
        with tf.variable_scope('layer_3'):
            output=util.conv2d(output,7,128)
            output=tf.nn.relu(output)
            print('rescale7' + str(output.get_shape().as_list()))
        with tf.variable_scope('layer_4'):
            output=util.conv2d(output,7,64)
            output=tf.nn.relu(output)
            print('rescale7' + str(output.get_shape().as_list()))
    return output

def train():
    input=tf.placeholder(tf.float32,shape=[batch_size,32,32,1])
    encoder_output = encoder(input)
    round_output = util.round_smooth(encoder_output)
    output=decoder(round_output)
    distortion=tf.reduce_mean(tf.squared_difference(input,output))
    learning_rate=tf.placeholder(tf.float32,shape=[])
    Moment=tf.train.MomentumOptimizer(learning_rate,0.9)
    train_op=Moment.minimize(distortion)
    init_op=tf.global_variables_initializer()
    with tf.Session() as sess:
        sess.run(init_op)
        data=util.data_iterator(batch_size)
        for i_iter in range(30000):
            if i_iter<10000:
                lr=0.05
            elif i_iter<20000:
                lr=0.01
            else:
                lr=0.001
            _,mse=sess.run([train_op,distortion],feed_dict={learning_rate:lr,input:next(data)})
            if i_iter % 1000==0:
                print('iter:{:>5d learning rate:{:>6.5f}=>{}'.format(i_iter,lr,mse))

if __name__=="__main__":
    train()





